from sqlalchemy import Column, String
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker

from init.mysql import engine

Base = declarative_base()


class CompanyIndex(Base):
    __tablename__ = 'company_index'
    code = Column(String(6))
    mc = Column(String(20), primary_key=True)
    dm = Column(String(20))
    ind = Column(String(10))
    outd = Column(String(10))

    def __init__(self,code, mc, dm, ind, outd):
        self.code = code
        self.mc = mc
        self.dm = dm
        self.ind = ind
        self.outd = outd

    @staticmethod
    def insert_or_update(code,data_list):
        session = sessionmaker(bind=engine)()
        for data in data_list:
            existing_company = session.query(CompanyIndex).filter_by(mc=data['mc'], code=code).first()
            if existing_company:
                existing_company.code = code
                existing_company.dm = data['dm']
                existing_company.ind = data['ind']
                existing_company.outd = data['outd']
            else:
                company_index = CompanyIndex(code,data['mc'], data['dm'], data['ind'], data['outd'])
                session.add(company_index)
        session.commit()


Base.metadata.create_all(engine)

